// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <exception>
#include <string>

#include <melon/exception.h>
#include <melon/fixed_string.h>
#include <melon/string.h>
#include <melon/synchronization/call_once.h>
#include <turbo/flags/flag.h>
#include <turbo/log/logging.h>
#include <pollux/flag_definitions/flags.h>
#include <pollux/common/process/stack_trace.h>

namespace kumo {
    namespace pollux {
        namespace error_source {
            using namespace melon::string_literals;

            /// Errors where the root cause of the problem is either because of bad input
            /// or an unsupported pattern of use are classified with source USER. Examples
            /// of errors in this category include syntax errors, unavailable names or
            /// objects.
            inline constexpr auto kErrorSourceUser = "USER"_fs;

            /// Errors where the root cause of the problem is an unexpected internal state
            /// in the system.
            inline constexpr auto kErrorSourceRuntime = "RUNTIME"_fs;

            /// Errors where the root cause of the problem is some unreliable aspect of the
            /// system are classified with source SYSTEM.
            inline constexpr auto kErrorSourceSystem = "SYSTEM"_fs;
        } // namespace error_source

        namespace error_code {
            using namespace melon::string_literals;

            ///====================== User Error Codes ======================:

            /// A generic user error code
            inline constexpr auto kGenericUserError = "GENERIC_USER_ERROR"_fs;

            /// An error raised when an argument verification fails
            inline constexpr auto kInvalidArgument = "INVALID_ARGUMENT"_fs;

            /// An error raised when a requested operation is not supported.
            inline constexpr auto kUnsupported = "UNSUPPORTED"_fs;

            /// Arithmetic errors - underflow, overflow, divide by zero etc.
            inline constexpr auto kArithmeticError = "ARITHMETIC_ERROR"_fs;

            /// An error raised when types are not compatible
            inline constexpr auto kSchemaMismatch = "SCHEMA_MISMATCH"_fs;

            ///====================== Runtime Error Codes ======================:

            /// An error raised when the current state of a component is invalid.
            inline constexpr auto kInvalidState = "INVALID_STATE"_fs;

            /// An error raised when unreachable code point was executed.
            inline constexpr auto kUnreachableCode = "UNREACHABLE_CODE"_fs;

            /// An error raised when a requested operation is not yet supported.
            inline constexpr auto kNotImplemented = "NOT_IMPLEMENTED"_fs;

            /// An error raised when memory pool exceeds limits.
            inline constexpr auto kMemCapExceeded = "MEM_CAP_EXCEEDED"_fs;

            /// An error raised when memory request failed due to arbitration failures. This
            /// is normally caused by insufficient global memory resource.
            inline constexpr auto kMemArbitrationFailure = "MEM_ARBITRATION_FAILURE"_fs;

            /// An error raised when memory pool is aborted.
            inline constexpr auto kMemAborted = "MEM_ABORTED"_fs;

            /// An error raised when memory arbitration times out.
            inline constexpr auto kMemArbitrationTimeout = "MEM_ARBITRATION_TIMEOUT"_fs;

            /// Error caused by memory allocation failure (inclusive of allocator memory cap
            /// exceeded).
            inline constexpr auto kMemAllocError = "MEM_ALLOC_ERROR"_fs;

            /// Error caused by failing to allocate cache buffer space for IO.
            inline constexpr auto kNoCacheSpace = "NO_CACHE_SPACE"_fs;

            /// An error raised when spill bytes exceeds limits.
            inline constexpr auto kSpillLimitExceeded = "SPILL_LIMIT_EXCEEDED"_fs;

            /// An error raised to indicate any general failure happened during spilling.
            inline constexpr auto kGenericSpillFailure = "GENERIC_SPILL_FAILURE"_fs;

            /// An error raised when trace bytes exceeds limits.
            inline constexpr auto kTraceLimitExceeded = "TRACE_LIMIT_EXCEEDED"_fs;

            /// Errors indicating file read corruptions.
            inline constexpr auto kFileCorruption = "FILE_CORRUPTION"_fs;

            /// Errors indicating file not found.
            inline constexpr auto kFileNotFound = "FILE_NOT_FOUND"_fs;

            /// We do not know how to classify it yet.
            inline constexpr auto kUnknown = "UNKNOWN"_fs;

            /// PolluxRuntimeErrors due to unsupported input values such as unicode input to
            /// cast-varchar-to-integer. This kind of errors is allowed in expression
            /// fuzzer.
            inline constexpr auto kUnsupportedInputUncatchable =
                    "UNSUPPORTED_INPUT_UNCATCHABLE"_fs;
        } // namespace error_code

        class PolluxException : public std::exception {
        public:
            enum class Type { kUser = 0, kSystem = 1 };

            PolluxException(
                const char *file,
                size_t line,
                const char *function,
                std::string_view expression,
                std::string_view message,
                std::string_view errorSource,
                std::string_view errorCode,
                bool isRetriable,
                Type exceptionType = Type::kSystem,
                std::string_view exceptionName = "PolluxException");

            /// Wrap an std::exception.
            PolluxException(
                const std::exception_ptr &e,
                std::string_view message,
                std::string_view errorSource,
                std::string_view errorCode,
                bool isRetriable,
                Type exceptionType = Type::kSystem,
                std::string_view exceptionName = "PolluxException");

            PolluxException(
                const std::exception_ptr &e,
                std::string_view message,
                std::string_view errorSource,
                bool isRetriable,
                Type exceptionType = Type::kSystem,
                std::string_view exceptionName = "PolluxException")
                : PolluxException(
                    e,
                    message,
                    errorSource,
                    "",
                    isRetriable,
                    exceptionType,
                    exceptionName) {
            }

            /// Inherited
            const char *what() const noexcept override {
                return state_->what();
            }

            /// Introduced nonvirtuals
            const process::StackTrace *stackTrace() const {
                return state_->stackTrace.get();
            }

            const char *file() const {
                return state_->file;
            }

            size_t line() const {
                return state_->line;
            }

            const char *function() const {
                return state_->function;
            }

            const std::string &failingExpression() const {
                return state_->failingExpression;
            }

            const std::string &message() const {
                return state_->message;
            }

            const std::string &errorCode() const {
                return state_->errorCode;
            }

            const std::string &errorSource() const {
                return state_->errorSource;
            }

            Type exceptionType() const {
                return state_->exceptionType;
            }

            const std::string &exceptionName() const {
                return state_->exceptionName;
            }

            bool isRetriable() const {
                return state_->isRetriable;
            }

            bool isUserError() const {
                return state_->errorSource == error_source::kErrorSourceUser;
            }

            const std::string &context() const {
                return state_->context;
            }

            const std::string &additionalContext() const {
                return state_->additionalContext;
            }

            const std::exception_ptr &wrappedException() const {
                return state_->wrappedException;
            }

        private:
            struct State {
                std::unique_ptr<process::StackTrace> stackTrace;
                Type exceptionType = Type::kSystem;
                std::string exceptionName;
                const char *file = nullptr;
                size_t line = 0;
                const char *function = nullptr;
                std::string failingExpression;
                std::string message;
                std::string errorSource;
                std::string errorCode;
                // The current exception context.
                std::string context;
                // The top-level ancestor of the current exception context.
                std::string additionalContext;
                bool isRetriable;
                // The original std::exception.
                std::exception_ptr wrappedException;

                mutable melon::once_flag once;
                mutable std::string elaborateMessage;

                template<typename F>
                static std::shared_ptr<const State> make(Type exceptionType, F);

                template<typename F>
                static std::shared_ptr<const State> make(F f) {
                    auto state = std::make_shared<PolluxException::State>();
                    f(*state);
                    return state;
                }

                void finalize() const;

                const char *what() const noexcept;
            };

            explicit PolluxException(std::shared_ptr<State const> state) noexcept
                : state_(std::move(state)) {
            }

            const std::shared_ptr<const State> state_;
        };

        class PolluxUserError : public PolluxException {
        public:
            PolluxUserError(
                const char *file,
                size_t line,
                const char *function,
                std::string_view expression,
                std::string_view message,
                std::string_view /* errorSource */,
                std::string_view errorCode,
                bool isRetriable,
                std::string_view exceptionName = "PolluxUserError")
                : PolluxException(
                    file,
                    line,
                    function,
                    expression,
                    message,
                    error_source::kErrorSourceUser,
                    errorCode,
                    isRetriable,
                    Type::kUser,
                    exceptionName) {
            }

            /// Wrap an std::exception.
            PolluxUserError(
                const std::exception_ptr &e,
                std::string_view message,
                bool isRetriable,
                std::string_view exceptionName = "PolluxUserError")
                : PolluxException(
                    e,
                    message,
                    error_source::kErrorSourceUser,
                    error_code::kInvalidArgument,
                    isRetriable,
                    Type::kUser,
                    exceptionName) {
            }
        };

        class PolluxRuntimeError final : public PolluxException {
        public:
            PolluxRuntimeError(
                const char *file,
                size_t line,
                const char *function,
                std::string_view expression,
                std::string_view message,
                std::string_view /* errorSource */,
                std::string_view errorCode,
                bool isRetriable,
                std::string_view exceptionName = "PolluxRuntimeError")
                : PolluxException(
                    file,
                    line,
                    function,
                    expression,
                    message,
                    error_source::kErrorSourceRuntime,
                    errorCode,
                    isRetriable,
                    Type::kSystem,
                    exceptionName) {
            }

            /// Wrap an std::exception.
            PolluxRuntimeError(
                const std::exception_ptr &e,
                std::string_view message,
                bool isRetriable,
                std::string_view exceptionName = "PolluxRuntimeError")
                : PolluxException(
                    e,
                    message,
                    error_source::kErrorSourceRuntime,
                    isRetriable,
                    Type::kSystem,
                    exceptionName) {
            }
        };

        /// Returns a reference to a thread level counter of Pollux error throws.
        int64_t &threadNumPolluxThrow();

        /// Returns a reference to a thread level boolean that controls whether no-throw
        /// APIs include detailed error messages in Status.
        bool &threadSkipErrorDetails();

        class ScopedThreadSkipErrorDetails {
        public:
            ScopedThreadSkipErrorDetails(bool skip = true)
                : original_{threadSkipErrorDetails()} {
                threadSkipErrorDetails() = skip;
            }

            ~ScopedThreadSkipErrorDetails() {
                threadSkipErrorDetails() = original_;
            }

        private:
            bool original_;
        };

        /// Holds a pointer to a function that provides addition context to be
        /// added to the detailed error message in case of an exception.
        struct ExceptionContext {
            using MessageFunction =
            std::string (*)(PolluxException::Type exceptionType, void *arg);

            /// Function to call in case of an exception to get additional context.
            MessageFunction messageFunc{nullptr};

            /// Value to pass to `messageFunc`. Can be null.
            void *arg{nullptr};

            /// If true, then the addition context in 'this' is always included when there
            /// are hierarchical exception contexts.
            bool isEssential{false};

            /// Pointer to the parent context when there are hierarchical exception
            /// contexts.
            ExceptionContext *parent{nullptr};

            /// Calls `messageFunc(arg)` and returns the result. Returns empty string if
            /// `messageFunc` is null.
            std::string message(PolluxException::Type exceptionType) {
                if (!messageFunc || suspended) {
                    return "";
                }

                std::string theMessage;

                try {
                    // Make sure not to call messageFunc again in case it throws.
                    suspended = true;
                    theMessage = messageFunc(exceptionType, arg);
                    suspended = false;
                } catch (...) {
                    return "Failed to produce additional context.";
                }

                return theMessage;
            }

            bool suspended{false};
        };

        /// If exceptionPtr represents an std::exception, convert it to PolluxUserError
        /// to add useful context for debugging.
        std::exception_ptr toPolluxException(const std::exception_ptr &exceptionPtr);

        /// Returns a reference to thread_local variable that holds a function that can
        /// be used to get addition context to be added to the detailed error message in
        /// case an exception occurs. This is to used in cases when stack trace would
        /// not provide enough information, e.g. in case of hierarchical processing like
        /// expression evaluation.
        ExceptionContext &getExceptionContext();

        /// RAII class to set and restore context for exceptions. Links the new
        /// exception context with the previous context held by the thread_local
        /// variable to allow retrieving the top-level context when there is an
        /// exception context hierarchy.
        class ExceptionContextSetter {
        public:
            explicit ExceptionContextSetter(ExceptionContext value)
                : prev_{getExceptionContext()} {
                value.parent = &prev_;
                getExceptionContext() = std::move(value);
            }

            ~ExceptionContextSetter() {
                getExceptionContext() = std::move(prev_);
            }

        private:
            ExceptionContext prev_;
        };
    } // namespace pollux
} // namespace kumo
